import vanna as vn

# # # 设置 API 密钥（可选）
# vn.set_api_key('your-api-key')

# 连接到数据库
vn.connect_to_postgres(
    host='localhost',
    dbname='test_db',
    user='root',
    password='123456',
    port=3306
)

# 训练模型
vn.train(ddl="CREATE TABLE sales (id INT, amount DECIMAL, date DATE)")

# 自然语言查询
result = vn.ask("显示最近一周的销售总额")
print(result)